package com.base.component.ai.agent.core;

import java.lang.reflect.InvocationTargetException;
import java.lang.reflect.Method;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.SortedMap;
import java.util.TreeMap;
import java.util.concurrent.CompletableFuture;
import java.util.function.Consumer;
import java.util.stream.Collectors;

import org.slf4j.MDC;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;

import com.base.component.ai.agent.bo.ContextBO;
import com.base.component.ai.agent.bo.ContextParamBO;
import com.base.component.ai.agent.core.base.Agent;
import com.base.core.context.utils.LoggerWrapper;
import com.gitee.magic.core.exception.ApplicationException;

/**
 * 
 * @author start
 *
 */
@Component
public class AgentWorkflowContext<MSG,AP extends ContextParamBO,PP> {

    @Autowired
    private Map<String, Agent<MSG,AP,PP>> agentMap;
	@Autowired
	private ContentTemplate<MSG,AP,PP> contentTemplate;
	@Autowired
	private ResponseHandler responseHandler;

	public void execute(ContextBO<MSG,AP,PP> context,Consumer<Object> ws) {
		List<Agent<MSG,AP,PP>> agentList=new ArrayList<>(agentMap.values());
		SortedMap<Integer,List<Agent<MSG,AP,PP>>>agentGroups=new TreeMap<>(agentList.stream().collect(Collectors.groupingBy(Agent::group)));
		executeAgent(context ,new ArrayList<>(agentGroups.values()),ws);
	}
	
	public void execute(ContextBO<MSG,AP,PP> context,List<List<String>> agentGroups,Consumer<Object> ws) {
		List<List<Agent<MSG,AP,PP>>> agentList=new ArrayList<>();
		for(List<String> agents : agentGroups) {
            agentList.add(agents.stream().filter(a->agentMap.containsKey(a)).map(a -> agentMap.get(a)).toList());
		}
		executeAgent(context ,agentList,ws);
	}

	private void executeAgent(ContextBO<MSG,AP,PP> context,List<List<Agent<MSG,AP,PP>>> agentGroups,Consumer<Object> ws) {
		AgentProxyService<MSG,AP,PP> chatProxyService=contentTemplate.getProxy(context.getBizType());
		context.setFrameConsumer((s)->{
			Class<?> responseCls=responseHandler.getClass();
			Consumer<Object> callbackResponse=(Object content)->{
				if(chatProxyService.msgHandle(context,content)){
					ws.accept(content);
				}
			};
			Boolean isOut=false;
			for(Method method:responseCls.getMethods()) {
				Object[] params=new Object[method.getParameterCount()];
				Boolean isInvoke=false;
				for(int i=0;i<method.getParameterCount();i++) {
					Class<?> cls=method.getParameterTypes()[i];
					if(cls.equals(ContextBO.class)) {
						params[i]=context;
					}else if(cls.equals(Consumer.class)) {
						params[i]=callbackResponse;
					}else if(cls.equals(s.getClass())) {
						params[i]=s;
						isInvoke=true;
					}
				}
				if(!isInvoke) {
					continue;
				}
				try {
					Object res=method.invoke(responseHandler,params);
					if(res!=null) {
						ws.accept(res);
					}
				} catch (IllegalAccessException e) {
					throw new ApplicationException(e);
				} catch (IllegalArgumentException e) {
					throw new ApplicationException(e);
				} catch (InvocationTargetException e) {
					if(e.getCause()!=null) {
						LoggerWrapper.print(e.getCause());
					}else {
						throw new ApplicationException(e);
					}
				}
				isOut=true;
				break;
			}
			if(!isOut) {
				callbackResponse.accept(s);
			}
		});
		//开始
		if(!chatProxyService.afterPropertiesSet(context,ws)) {
			return;
		}
		try {
			Map<String,String> mdcMap=MDC.getCopyOfContextMap();
			for(List<Agent<MSG,AP,PP>> agentList : agentGroups) {
//	            List<Agent<MSG,AP,PP>> agentList = new ArrayList<>(
//	            		agents.stream().filter(a->agentMap.containsKey(a)).map(a -> agentMap.get(a)).toList());
				if(context.isTerminate()) {
					//终止
					break;
				}
	            List<CompletableFuture<Void>> agentFutures = agentList.stream()
	            		.filter(t->t.support(context.getMsg().getClass()))
	            		.map(a -> CompletableFuture.runAsync(() -> {
	            		MDC.setContextMap(mdcMap);
	    			try {
	    				a.apply(context);
	    			}catch (Exception e) {
	    				if(a.exception(e,context)) {
	    					LoggerWrapper.print(e);
	    				}
	    			}
	            })).toList();
	            CompletableFuture.allOf(agentFutures.toArray(new CompletableFuture[0])).join();
	        }
		}catch(Exception e) {
			LoggerWrapper.print(e);
		}finally {
			 chatProxyService.destory(context);
		}
	}
	
}
